import os
import platform
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from datasets import load_from_disk
from transformers import AdamW
from transformers import T5Tokenizer, T5ForConditionalGeneration
# from logging_util import get_logger
from rouge import Rouge
from transformers.utils import ExplicitEnum
import MengZiT5Model

device = 'cuda' if torch.cuda.is_available() else 'cpu'

# 获取当前操作系统的名称
os_name = platform.system()
# logger = get_logger(model_name='mengzi-t5-base')

# 设置模型路径及数据集路径
if os_name == "Windows":
    model_dir = r'D:\python\models\langboat\meng_zi_t5'
    data_dir = r'D:\python\datas\nlp_seq2seq\nlpcc_2017'
    print("当前执行环境是 Windows...")
elif os_name == "Linux":
    model_dir = r'/root/autodl-fs/models/meng_zi_t5'
    data_dir = r'/root/autodl-fs/data/nlp_ai/nlp_seq2seq/nlpcc_2017'
    print("当前执行环境是 Linux...")
else:
    raise ValueError("当前执行环境不是 Windows 也不是 Linux")

class MengZiT5Model(nn.Module):
    def __init__(self):
        super().__init__()
        # 加载预训练模型
        self.model = T5ForConditionalGeneration.from_pretrained(model_dir)


    def forward(self, inputs, labels=None):
        # 1、encoder的input_ids和attention_mask
        input_ids = inputs['input_ids']
        attention_mask = inputs['attention_mask']

        if labels is not None:
            # 2、decoder 的labels
            train_labels = labels['input_ids'].contiguous()
            train_labels_mask = labels['attention_mask']

            # 3、decoder 的input_ids和attention_mask
            decoder_input_ids = train_labels.new_zeros(train_labels.shape)
            decoder_input_ids[..., 1:] = train_labels[..., :-1].clone()

            decoder_attention_mask = train_labels_mask.new_zeros(train_labels_mask.shape)
            decoder_attention_mask[..., 1:] = train_labels_mask[..., :-1].clone()
            decoder_attention_mask[..., 0] = 1
            # 4、送入模型进行预测
            outputs = self.model(input_ids=input_ids
                                 , attention_mask=attention_mask
                                 , decoder_input_ids=decoder_input_ids
                                 , decoder_attention_mask=decoder_attention_mask
                                 , labels=train_labels)
            # 5、返回训练时候的Loss值
            return outputs.loss
        else:
            # 模型生成
            summary_ids = self.model.generate(input_ids
                                              , num_beams=4 # 束搜索法
                                              , no_repeat_ngram_size=2 # 确保不重复
                                              , min_length=10 # 长度限制
                                              , max_length=64
                                              , early_stopping=True
            )
            # 将id转换为输出 summary_ids.shape = [bs, length]
            outputs = tokenizer.batch_decode(summary_ids, skip_special_tokens=True)
            return outputs
